package com.bestv.search.engine.compass.extend;

import org.apache.lucene.search.ScoreDoc;
import org.apache.lucene.search.ScoreDocComparator;
import org.apache.lucene.search.SortField;

import java.util.Map;

public class TagsMaxHitCountComparator implements ScoreDocComparator {
    private String[] tagsInHd;
    private Map<String, Integer> maxHitCountByTagsMap;

    public TagsMaxHitCountComparator(String[] tagsInHd, Map<String, Integer> maxHitCountByTagsMap) {
        this.tagsInHd = tagsInHd;
        this.maxHitCountByTagsMap = maxHitCountByTagsMap;
    }

    @Override
    public int compare(ScoreDoc doc1, ScoreDoc doc2) {
        String tag1 = tagsInHd[doc1.doc];
        Integer maxHitCount1 = maxHitCountByTagsMap.get(tag1);
        String tag2 = tagsInHd[doc2.doc];
        Integer maxHitCount2 = maxHitCountByTagsMap.get(tag2);

        maxHitCount1 = maxHitCount1 == null ? 0 : maxHitCount1;
        maxHitCount2 = maxHitCount2 == null ? 0 : maxHitCount2;

        // Order by maxHitCount desc
        if (maxHitCount1.intValue() > maxHitCount2.intValue()) {
            return -1;
        } else if (maxHitCount1.intValue() == maxHitCount2.intValue()) {
            return 0;
        } else {
            return 1;
        }
    }

    @Override
    public int sortType() {
        return SortField.INT;
    }

    @Override
    public Comparable<Integer> sortValue(ScoreDoc doc) {
        String tag = tagsInHd[doc.doc];
        Integer maxHitCount = maxHitCountByTagsMap.get(tag);
        if (maxHitCount == null) {
            return 0;
        } else {
            return maxHitCount.intValue();
        }
    }
}